{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment this cell if running in Google Colab\n",
    "# !pip install clinicadl==0.0.3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train your own models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This section explains how to train a CNN on OASIS data that were processed in the previous sections. \n",
    "\n",
    "```{warning}\n",
    "If you do not have access to a GPU, training the CNN may require too much time. However, you can execute this notebook on Colab to run it on a GPU.\n",
    "```\n",
    "\n",
    "If you already know the models implemented in `clinicadl`, you can directly jump to the [last section](./training.ipynb#Customize-your-experiment!) to implement your own custom experiment!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "# Check if a GPU is available\n",
    "print('GPU is available', torch.cuda.is_available())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-warning\">\n",
    "<b>Data used for training:</b><p>\n",
    "    Because they are time-costly, the preprocessing steps presented in the beginning of this tutorial were only executed on a subset of OASIS-1, but obviously two participants are insufficient to train a network! To obtain more meaningful results, you should retrieve the whole <a href=\"https://www.oasis-brains.org/\">OASIS-1</a> dataset and run the training based on the labels and splits performed in the previous section.</p>\n",
    "    <p>Of course, you can use another dataset, but then you will have to perform again <a href=\"./label_extraction.ipynb\">the extraction of labels and data splits</a> on this dataset.</p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "### Input types\n",
    "\n",
    "In this notebook the goal is to train a classifier to differentiate the **AD label from the CN label** based on the associated **T1-MR images**. These images were preprocessed according to the section on [preprocessing](./preprocessing.ipynb).\n",
    "\n",
    "According to the literature review done in ([Wen et al., 2020](https://www.sciencedirect.com/science/article/abs/pii/S1361841520300591)), 4 main paradigms were implemented to train a CNN to perform a classification task between AD and CN labels based on T1-MRI. These paradigms depend of the input given to the network:\n",
    "* 2D slices,\n",
    "* 3D patches,\n",
    "* ROIs (3D patches focused on a region of interest),\n",
    "* 3D images.\n",
    "\n",
    "### Architectures\n",
    "\n",
    "There is no consensus in the literature on the architecture of the models that should be used for each input category. Part of the studies reviewed in ([Wen et al., 2020](https://www.sciencedirect.com/science/article/abs/pii/S1361841520300591)) used custom models, and the others reused architectures that led to good results on natural images (VGGNet, ResNet, GoogleNet, DenseNet...).\n",
    "\n",
    "We chose to find custom architectures by running a search on the main architecture components (more details can be found in section 4.3 of ([Wen et al., 2020](https://www.sciencedirect.com/science/article/abs/pii/S1361841520300591))). This architecture search was run only on the train + validation sets, and **no test sets were used to choose the best architectures**. This is crucial, as the use of a test set during architecture search can lead to biased, over-optimistic, results.\n",
    "\n",
    "### Training networks with `clinicadl`\n",
    "\n",
    "The training tasks can be performed using `clinicadl`:\n",
    "```bash\n",
    "clinicadl train <input_type> <network_type> <caps_directory> \\\n",
    "                <preprocessing> <tsv_path> <output_directory> <architecture>\n",
    "```\n",
    "where:\n",
    "- `input_type` is the type of input used. Must be chosen between `image`, `patch`, `roi` and `slice`.\n",
    "- `network_type` is the type of network trained. The possibilities depend on the `input_type`, but the complete list of options is `autoencoder`, `cnn` or `multicnn`.\n",
    "- `caps_directory` is the input folder containing the results of the [`t1-linear` pipeline](./preprocessing.ipynb#Image-preprocessing-with-the-t1-linear-pipeline).\n",
    "- `preprocessing` corresponds to the preprocessing pipeline whose outputs will be used for training. Current version only supports `t1-linear`. `t1-extensive` will be implemented in next versions of clinicadl.\n",
    "- `tsv_path` is the input folder of a tsv file tree generated by [`clinicadl tsvtool {split|kfold}`](./label_extraction.ipynb#Build-test-set).\n",
    "- `output_directory` is the folder where the results are stored.\n",
    "- `architecture` is the name of the architecture used (for example, `Conv5_FC3`).\n",
    "\n",
    "\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "<b>Computational resources:</b><p>\n",
    "    Depending on your computational resources, you should adapt the parameters in the <b>Computational resources</b> argument group:\n",
    "    <ul>\n",
    "        <li> <code>-cpu</code> forces the system to use CPU instead of looking for a GPU. </li>\n",
    "        <li> <code>--nproc N</code> sets the number of workers (parallel processes) in the DataLoader to N. </li>\n",
    "        <li> <code>--batch_size B</code> will set the batch size to B. If the batch size is too high, it can raise memory errors. </li>\n",
    "    </ul>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2D slice-level CNN\n",
    "\n",
    "An advantage of the 2D slice-level approach is that existing CNNs which had huge success for natural image classification, e.g. ResNet ([He et al., 2016](https://doi.org/10.1109/CVPR.2016.90)) and VGGNet ([Simonyan and Zisserman, 2014](https://arxiv.org/abs/1409.1556)) , can be easily borrowed and used in a transfer learning fashion. Other advantages are the increased number of training samples as many slices can be extracted from a single 3D image, and a lower memory usage compared to using the full MR image as input.\n",
    "\n",
    "The 2D slice-level CNN in `clinicadl` is a slight modification of the ResNet-18 network used to train natural images on the ImageNet dataset:\n",
    "* One fully-connected layer was added at the end of the network to reduce the dimension of the output from 1000 to 2 classes <font color='purple'>(purple dotted box)</font>.\n",
    "* The last five convolutional layers and the last FC of ResNet are fine-tuned <font color='green'>(green dotted box)</font>.\n",
    "* All other layers have their weights and biases fixed.\n",
    "\n",
    "<img src=\"./images/2DCNN.png\">\n",
    "\n",
    "During training, the gradients update are done based on the loss computed at the slice level. Final performance metric are computed at the subject level by combining the outputs of the slices of the same subject.\n",
    "\n",
    "<div class=\"alert alert-block alert-success\">\n",
    "<b>Transfer learning:</b> \n",
    "    <p>The layers of this network (except the last one) are initialized with the weights of the ResNet-18 trained on ImageNet.</p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!clinicadl train slice -h\n",
    "#!clinicadl train slice <caps_dir> \"t1-linear\" data/labels_list results/slice_cnn resnet18"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3D patch-level & ROI-based models\n",
    "\n",
    "The 3D patch-level models compensate the absence of 3D information in the 2D slice-level approach and keep some of its advantages (low memory usage and larger sample size). This paradigm can be divided into two different frameworks:\n",
    "- **single-CNN**: one CNN is trained on all patch locations.\n",
    "- **multi-CNN**: one CNN is trained per patch location.\n",
    "For **multi-CNN** the sample size is smaller (equivalent to image level framework), however the CNNs may be more accurate as they are specialized for one patch location.\n",
    "\n",
    "ROI-based models are similar to 3D-patch but take advantage of prior knowledge on Alzheimer's disease. Indeed most of the patches are not informative as they contain parts of the brain that are not affected by the disease. Methods based on regions of interest (ROI) overcome this issue by focusing on regions which are known to be informative: here the hippocampi. In this way, the complexity of the framework can be decreased as fewer inputs are used to train the networks.\n",
    "\n",
    "The 3D patch-level and ROI-based CNNs in `clinicadl` have the same architecture, including:\n",
    "* 4 convolutional layers with kernel 3x3x3,\n",
    "* 4 max pooling layers with stride and kernel of 2 and a padding value that automatically adapts to the input feature map size.\n",
    "* 3 fully-connected layers.\n",
    "\n",
    "<img src=\"./images/ROICNN.png\">\n",
    "\n",
    "As for the 2D slice-level model, the gradient updates are done based on the loss computed at the patch level. Final performance metrics are computed at the subject level by combining the outputs of the patches or the two hippocampi of the same subject.\n",
    "\n",
    "<div class=\"alert alert-block alert-success\">\n",
    "<b>Transfer learning:</b> \n",
    "    <p>It is possible for these categories to train an autoencoder derived from the CNN architecture. The encoder will share the same architecture as the CNN until the fully-connected layers (see the <a href=\"./deep_learning.ipynb#Autoencoder-pretraining\">bakground section</a> for more details on autoencoders construction).</p>\n",
    "    <img src=\"./images/autoencoder.png\">\n",
    "    <p>Then the weights of the encoder will be transferred to the convolutions of the CNN to initialize it before its training. This procedure is called <i>autoencoder pretraining</i>.</p>\n",
    "    <p>It is also possible to transfer weights between two CNNs with the same architecture.</p>\n",
    "    <p>For 3D-patch multi-CNNs specifically, it is possible to initialize each CNN of a multi-CNN:\n",
    "        <ul>\n",
    "        <li> with the weights of a single-CNN,</li>\n",
    "        <li> with the weights of the corresponding CNN of a multi-CNN. </li>\n",
    "    </ul>\n",
    "    <p>Transferring weights between CNNs can be useful when performing two classification tasks that are similar. This is what has been done in (<a href=\"https://www.sciencedirect.com/science/article/abs/pii/S1361841520300591\">Wen et al., 2020</a>): the sMCI vs pMCI classification network was initialized with the weights of the AD vs CN classification network.</p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{warning}\n",
    "Transferring weights between tasks that are not similar enough can hurt the performance!\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3D-patch level models\n",
    "\n",
    "See definition of patches in [the background section on neuroimaging]()."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 3D-patch autoencoder pretraining\n",
    "!clinicadl train patch autoencoder -h\n",
    "#!clinicadl train patch autoencoder <caps_dir> \"t1-linear\" data/labels_list results/patch_autoencoder Conv4_FC3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 3D-patch single-CNN training\n",
    "!clinicadl train patch cnn -h\n",
    "\n",
    "# With autoencoder pretraining\n",
    "#!clinicadl train patch cnn <caps_dir> \"t1-linear\" data/labels_list results/patch_single-cnn Conv4_FC3 --transfer_learning_path results/patch_autoencoder\n",
    "# Without pretraining\n",
    "#!clinicadl train patch cnn <caps_dir> \"t1-linear\" data/labels_list results/patch_single-cnn Conv4_FC3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 3D-patch multi-CNN training\n",
    "!clinicadl train patch multicnn -h\n",
    "\n",
    "# With autoencoder pretraining\n",
    "#!clinicadl train patch multicnn <caps_dir> \"t1-linear\" data/labels_list results/patch_multi-cnn Conv4_FC3 --transfer_learning_path results/patch_autoencoder\n",
    "# With single-CNN pretraining\n",
    "#!clinicadl train patch multicnn <caps_dir> \"t1-linear\" data/labels_list results/patch_multi-cnn Conv4_FC3 --transfer_learning_path results/patch_single-cnn\n",
    "# Without pretraining\n",
    "#!clinicadl train patch multicnn <caps_dir> \"t1-linear\" data/labels_list results/patch_multi-cnn Conv4_FC3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ROI-based models\n",
    "\n",
    "ROI inputs correspond to two patches of size 50x50x50 manually centered on each hippocampus. They cannot be pre extracted with `clinicadl extract` and are computed from the whole MR volume.\n",
    "\n",
    "<p align=\"center\"><img src=\"./images/hippocampi.png\" alt=\"Coronal views of ROI inputs\"></p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ROI-based autoencoder pretraining\n",
    "!clinicadl train roi autoencoder -h\n",
    "#!clinicadl train roi autoencoder <caps_dir> \"t1-linear\" data/labels_list results/roi_autoencoder Conv4_FC3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ROI-based CNN training\n",
    "!clinicadl train roi cnn -h\n",
    "\n",
    "# With autoencoder pretraining\n",
    "#!clinicadl train roi cnn <caps_dir> \"t1-linear\" data/labels_list results/roi_cnn Conv4_FC3 --transfer_learning_path results/roi_autoencoder\n",
    "# Without pretraining\n",
    "#!clinicadl train roi cnn <caps_dir> \"t1-linear\" data/labels_list results/roi_cnn Conv4_FC3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3D-image level model\n",
    "\n",
    "In this approach, the whole MRI is used at once and the classification is performed at the image level. The advantage is that the spatial information is fully integrated, and it may allow the discovery of new knowledge on the disease. However, it requires more computational resources (especially GPU with higher memory capacity).\n",
    "\n",
    "The 3D image-level CNN in `clinicadl` is designed as follows:\n",
    "* 5 convolutional layers with kernel 3x3x3,\n",
    "* 5 max pooling layers with stride and kernel of 2 and a padding value that automatically adapts to the input feature map size.\n",
    "* 3 fully-connected layers.\n",
    "\n",
    "<img src=\"./images/imageCNN.png\">\n",
    "\n",
    "Depending on the preprocessing, the size of the fully connected layers must be adapted. This is why two models exist in `clinicadl`:\n",
    "* `Conv5_FC3` adapted to `t1-linear` preprocessing,\n",
    "* `Conv5_FC3_mni` adapted to `t1-extensive` preprocessing (coming soon in next releases of clinicadl).\n",
    "\n",
    "\n",
    "<div class=\"alert alert-block alert-success\">\n",
    "<b>Transfer learning:</b> \n",
    "    <p>It is possible for this category to train an autoencoder derived from the CNN architecture, or to transfer weights between CNNs. See the section on patches for more details on this topic!</p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 3D-image level autoencoder pretraining\n",
    "!clinicadl train image autoencoder -h\n",
    "#!clinicadl train image autoencoder <caps_dir> \"t1-linear\" data/labels_list results/image_autoencoder Conv5_FC3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 3D-image level autoencoder pretraining\n",
    "!clinicadl train image cnn -h\n",
    "\n",
    "# With autoencoder pretraining\n",
    "#!clinicadl train image cnn <caps_dir> \"t1-linear\" data/labels_list results/image_cnn Conv5_FC3 --transfer_learning_path results/image_autoencoder\n",
    "# Without pretraining\n",
    "#!clinicadl train image cnn <caps_dir> \"t1-linear\" data/labels_list results/image_cnn Conv5_FC3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Customize your experiment!\n",
    "\n",
    "You want to train your custom architecture, with a custom input type or preprocessing on other labels? Please fork and clone the [github repo](https://github.com/aramis-lab/AD-DL]) to add your changes.\n",
    "\n",
    "\n",
    "### Add a custom model <i class=\"fa fa-hourglass-start \" aria-hidden=\"true\"></i>\n",
    "Write your model class in `clinicadl/tools/deep_learning/models` and import it in `clinicadl/tools/deep_learning/models/__init__.py`.\n",
    "\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "<b>Autoencoder transformation:</b><p>\n",
    "    Your custom model can be transformed in autoencoder in the same way as predefined models. To make it possible, implement the convolutional part in <code>features</code> and the fully-connected layer in <code>classifier</code>. See predefined models as examples.\n",
    "</div>\n",
    "\n",
    "### Add a custom input type <i class=\"fa fa-hourglass-start \" aria-hidden=\"true\"></i> <i class=\"fa fa-hourglass-start \" aria-hidden=\"true\"></i> <i class=\"fa fa-hourglass-start \" aria-hidden=\"true\"></i>\n",
    "\n",
    "Input types that are already provided in `clinicadl` are image, patch, roi and slice. To add a custom input type, please follow the steps detailed below:\n",
    "* Choose a mode name for this input type (for example default ones are image, patch, roi and slice). \n",
    "* Add your dataset class in `clinicadl/tools/deep_learning/data.py` as a child class of the abstract class `MRIDataset`.\n",
    "* Create your dataset in `return_dataset` by adding:\n",
    "```\n",
    "elif mode==<mode_name>:\n",
    "    return <dataset_class>(\n",
    "        input_dir,\n",
    "        data_df,\n",
    "        preprocessing=preprocessing,\n",
    "        transformations=transformations,\n",
    "        <custom_args>\n",
    "    )\n",
    "```\n",
    "* Add your custom subparser to `train` and complete `train_func` in `clinicadl/cli.py`.\n",
    "\n",
    "### Add a custom preprocessing <i class=\"fa fa-hourglass-start \" aria-hidden=\"true\"></i> <i class=\"fa fa-hourglass-start \" aria-hidden=\"true\"></i>\n",
    "Define the path of your new preprocessing in the `_get_path` method of `MRIDataset` in `clinicadl/tools/deep_learning/data.py`. You will also have to add the name of your preprocessing pipeline in the general command line by modifying the possible choices of the `preprocessing` argument of `train_pos_group` in `cli.py`.\n",
    "\n",
    "### Change the labels <i class=\"fa fa-hourglass-start \" aria-hidden=\"true\"></i>\n",
    "You can launch a classification task with clinicadl using any labels. The input tsv files must include the columns `participant_id`, `session_id` and `diagnosis`. If the column `diagnosis` does not contain the labels described in this tutorial (AD, CN, MCI, sMCI, pMCI), you can add your own label name associated to a class value in the `diagnosis_code` of the class `MRIDataset` in `clinicadl/tools/deep_learning/data.py`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{admonition} Suggestion!\n",
    ":class: tip\n",
    "Do not hesitate to ask for help on [GitHub](https://github.com/aramis-lab/AD-DL/issues/new) or propose a new pull request!\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
